/// by Eric sun 2025-08-31

#include "filter.h"
#include "ioctl.h"

// 简单的用户/内核通信实现
// 支持：
//  - IOCTL_WDDM_FILTER_ADD   (结构化输入 WDDM_FILTER_ADD_INPUT -> 返回 ULONG id)
//  - IOCTL_WDDM_FILTER_REMOVE (输入 ULONG id，id==0 删除最后一个，输出 ULONG success)
//  - 回显（任何未匹配的 METHOD_BUFFERED 请求会被回显回输出缓冲区）

static LONG g_virtual_monitor_count = 0;

// Worker routine used to call unload from a different thread context
static VOID wddm_unload_worker(PDEVICE_OBJECT DeviceObject, PVOID Context)
{
	UNREFERENCED_PARAMETER(DeviceObject);
	PIO_WORKITEM work = (PIO_WORKITEM)Context;
	DPT("wddm_unload_worker: invoking wddm_filter_unload\n");
	if (wf && wf->driver_object) {
		wddm_filter_unload(wf->driver_object);
	}
	// Free the work item allocated by IoAllocateWorkItem
	if (work) {
		IoFreeWorkItem(work);
	}
}

// handle_ioctl: 处理来自用户态的 DeviceIoControl 请求（METHOD_BUFFERED 场景）
NTSTATUS handle_ioctl(PIRP irp, PIO_STACK_LOCATION irp_stack)
{
	NTSTATUS status = STATUS_SUCCESS;
	ULONG ioctl_code = irp_stack->Parameters.DeviceIoControl.IoControlCode;
	PVOID inputBuffer = irp->AssociatedIrp.SystemBuffer;
	PVOID outputBuffer = irp->AssociatedIrp.SystemBuffer;
	ULONG inputLength = irp_stack->Parameters.DeviceIoControl.InputBufferLength;
	ULONG outputLength = irp_stack->Parameters.DeviceIoControl.OutputBufferLength;

	DPT("handle_ioctl: code=0x%X, in=%u out=%u\n", ioctl_code, inputLength, outputLength);

	// 1) 结构化 ADD IOCTL
	if (ioctl_code == IOCTL_WDDM_FILTER_ADD && inputBuffer && inputLength >= sizeof(WDDM_FILTER_ADD_INPUT)) {
		PWDDM_FILTER_ADD_INPUT in = (PWDDM_FILTER_ADD_INPUT)inputBuffer;
		LONG new_id = InterlockedIncrement(&g_virtual_monitor_count);
		PVIRTUAL_CHILD_ENTRY entry = (PVIRTUAL_CHILD_ENTRY)ExAllocatePoolWithTag(NonPagedPool, sizeof(VIRTUAL_CHILD_ENTRY), 'vChd');
		if (!entry) {
			status = STATUS_INSUFFICIENT_RESOURCES;
			irp->IoStatus.Status = status;
			irp->IoStatus.Information = 0;
			IoCompleteRequest(irp, IO_NO_INCREMENT);
			return status;
		}

		RtlZeroMemory(entry, sizeof(*entry));
		entry->ChildUid = (ULONG)new_id;
		entry->Width = in->Width ? in->Width : 1024;
		entry->Height = in->Height ? in->Height : 768;
		// 复制名称并保证 NUL 末尾
		RtlStringCbCopyNA(entry->Name, sizeof(entry->Name), in->Name, sizeof(in->Name));
		entry->Name[sizeof(entry->Name)-1] = '\0';

		wf_lock();
		InsertTailList(&wf->virtual_child_head, &entry->ListEntry);
		InterlockedIncrement(&wf->virtual_child_count);
		wf_unlock();

		if (outputBuffer && outputLength >= sizeof(ULONG)) {
			*(ULONG*)outputBuffer = (ULONG)new_id;
			irp->IoStatus.Information = sizeof(ULONG);
			status = STATUS_SUCCESS;
		} else {
			status = STATUS_BUFFER_TOO_SMALL;
			irp->IoStatus.Information = 0;
		}

		irp->IoStatus.Status = status;
		IoCompleteRequest(irp, IO_NO_INCREMENT);
		return status;
	}

	// 2) 兼容老字符串命令："ADD_VMON"（保持向后兼容）
	if (inputBuffer && inputLength >= 8) {
		if (RtlCompareMemory(inputBuffer, "ADD_VMON", 8) == 8) {
			LONG new_id = InterlockedIncrement(&g_virtual_monitor_count);
			PVIRTUAL_CHILD_ENTRY entry = (PVIRTUAL_CHILD_ENTRY)ExAllocatePoolWithTag(NonPagedPool, sizeof(VIRTUAL_CHILD_ENTRY), 'vChd');
			if (!entry) {
				status = STATUS_INSUFFICIENT_RESOURCES;
				irp->IoStatus.Status = status;
				irp->IoStatus.Information = 0;
				IoCompleteRequest(irp, IO_NO_INCREMENT);
				return status;
			}

			RtlZeroMemory(entry, sizeof(*entry));
			entry->ChildUid = (ULONG)new_id;
			entry->Width = 1024;
			entry->Height = 768;
			RtlStringCbCopyA(entry->Name, sizeof(entry->Name), "VirtualDisplay");

			wf_lock();
			InsertTailList(&wf->virtual_child_head, &entry->ListEntry);
			InterlockedIncrement(&wf->virtual_child_count);
			wf_unlock();

			if (outputBuffer && outputLength >= sizeof(ULONG)) {
				*(ULONG*)outputBuffer = (ULONG)new_id;
				irp->IoStatus.Information = sizeof(ULONG);
				status = STATUS_SUCCESS;
			} else {
				status = STATUS_BUFFER_TOO_SMALL;
				irp->IoStatus.Information = 0;
			}

			irp->IoStatus.Status = status;
			IoCompleteRequest(irp, IO_NO_INCREMENT);
			return status;
		}
	}

	// 3) REMOVE IOCTL: 输入输出缓冲区为 ULONG id，id==0 表示删除最后一个
	if (ioctl_code == IOCTL_WDDM_FILTER_REMOVE && inputLength >= sizeof(ULONG) && outputLength >= sizeof(ULONG)) {
		ULONG id = *(ULONG*)inputBuffer;
		BOOLEAN removed = FALSE;

		wf_lock();
		PLIST_ENTRY p = wf->virtual_child_head.Flink;
		PVIRTUAL_CHILD_ENTRY last = NULL;
		while (p != &wf->virtual_child_head) {
			PVIRTUAL_CHILD_ENTRY v = CONTAINING_RECORD(p, VIRTUAL_CHILD_ENTRY, ListEntry);
			p = p->Flink;
			last = v;
			if (id != 0 && v->ChildUid == id) {
				RemoveEntryList(&v->ListEntry);
				InterlockedDecrement(&wf->virtual_child_count);
				removed = TRUE;
				ExFreePoolWithTag(v, 'vChd');
				break;
			}
		}
		if (!removed && id == 0 && last) {
			RemoveEntryList(&last->ListEntry);
			InterlockedDecrement(&wf->virtual_child_count);
			removed = TRUE;
			ExFreePoolWithTag(last, 'vChd');
		}
		wf_unlock();

		if (removed) {
			*(ULONG*)outputBuffer = 1;
			irp->IoStatus.Information = sizeof(ULONG);
			irp->IoStatus.Status = STATUS_SUCCESS;
			IoCompleteRequest(irp, IO_NO_INCREMENT);
			return STATUS_SUCCESS;
		} else {
			irp->IoStatus.Status = STATUS_NOT_FOUND;
			irp->IoStatus.Information = 0;
			IoCompleteRequest(irp, IO_NO_INCREMENT);
			return STATUS_NOT_FOUND;
		}
	}

	// 4) 通用回显：如果用户传入数据且输出缓冲区足够大，将输入直接拷贝回输出
	if (inputBuffer && outputBuffer && inputLength > 0 && outputLength >= inputLength) {
		RtlCopyMemory(outputBuffer, inputBuffer, inputLength);
		irp->IoStatus.Information = inputLength;
		irp->IoStatus.Status = STATUS_SUCCESS;
		IoCompleteRequest(irp, IO_NO_INCREMENT);
		return STATUS_SUCCESS;
	}

	// 5) 支持卸载 IOCTL（异步触发），安全地在工作线程中调用卸载例程
	if (ioctl_code == IOCTL_WDDM_FILTER_UNLOAD) {
		// 创建并排队一个工作项以在非 IRQL 高的上下文中执行卸载
		PDEVICE_OBJECT dev = wf ? wf->ctrl_devobj : NULL;
		PIO_WORKITEM work = NULL;
		if (dev) {
			work = IoAllocateWorkItem(dev);
		}
		if (work) {
			IoQueueWorkItem(work, (PIO_WORKITEM_ROUTINE)wddm_unload_worker, DelayedWorkQueue, work);
			// 我们不会在这里释放 work，因为卸载工作中会释放/忽略相应资源
			if (outputBuffer && outputLength >= sizeof(ULONG)) {
				*(ULONG*)outputBuffer = 1;
				irp->IoStatus.Information = sizeof(ULONG);
			}
			irp->IoStatus.Status = STATUS_SUCCESS;
			IoCompleteRequest(irp, IO_NO_INCREMENT);
			return STATUS_SUCCESS;
		} else {
			irp->IoStatus.Status = STATUS_INSUFFICIENT_RESOURCES;
			irp->IoStatus.Information = 0;
			IoCompleteRequest(irp, IO_NO_INCREMENT);
			return STATUS_INSUFFICIENT_RESOURCES;
		}
	}

	// 未知或不支持的 IOCTL：返回 STATUS_INVALID_DEVICE_REQUEST
	DPT("handle_ioctl: unsupported IOCTL 0x%X\n", ioctl_code);
	irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;
	irp->IoStatus.Information = 0;
	IoCompleteRequest(irp, IO_NO_INCREMENT);
	return STATUS_INVALID_DEVICE_REQUEST;
}


// 对外暴露的 IOCTL 处理入口，按驱动设计可能由 common dispatch 使用
NTSTATUS wddm_filter_handle_ioctl(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	PIO_STACK_LOCATION irpSp = IoGetCurrentIrpStackLocation(Irp);
	if (!irpSp) {
		Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
		Irp->IoStatus.Information = 0;
		IoCompleteRequest(Irp, IO_NO_INCREMENT);
		return STATUS_INVALID_PARAMETER;
	}

	// 仅处理 METHOD_BUFFERED 场景的用户 IOCTL（如果需要，可以扩展处理 METHOD_NEITHER/METHOD_IN_DIRECT 等）
	return handle_ioctl(Irp, irpSp);
}


